import numpy as np
import torch
from torch import nn


class AttentionMask(nn.Module):

    @staticmethod
    def get_padding_mask(inputs_q, len_k):
        """
        padding mask的作用：不同batch之间句子长度可以不一样，但是每个batch的长度必须是一样的：
        因此出现一个问题，不够长度需要加pad，使得其长度变成一样。
        :param inputs_q: shape(batch_size, context_length)
        :param len_k: context_length_k
        """
        # inputs只是给个维度
        batch_size, len_q = inputs_q.size()
        padding_mask = inputs_q.data.eq(0).unsqueeze(1)
        # shape(batch_size,1, len_q)
        # inputs.data.eq(0)，这句的作用是返回一个大小和seq_k一样的 tensor，只不过里面的值只有 True 和 False。
        # 如果 inputs 某个位置的值等于 0，那么对应位置就是True，否则即为False。
        # 举个例子，输入为 inputs = [1, 2, 3, 4, 0]，
        # inputs.data.eq(0)就会返回[False, False, False, False, True]

        return padding_mask.expand(batch_size, len_q, len_k)

    ''' 因为enc是处理K,V的，求Attention时是用v1,v2,..vm去加权的，要把pad对应的v_i的相关系数设为0，这样注意力就不会关注pad向量

    enc_self_attn_mask = get_attn_pad_mask(enc_inputs, enc_inputs) = enc_inputs自注意力层的时候PAD部分
            tensor([[[False, False, False, False,  True],
                    [False, False, False, False,  True],     横轴enc_inputs(ich mochte ein bier P) 和 竖轴enc_outputs(ich mochte ein bier P)
                    [False, False, False, False,  True],     竖轴经注意力机制后生成得ich是带有原ich mochte ein bier P关系得ich 它和原P没关系 原P仅仅是填充
                    [False, False, False, False,  True],     最后一行P也是[False, False, False, False,  True]因为 这个新得填充P也要和原语句建立联系
                    [False, False, False, False,  True]]])         当对每一行做Softmax时候，我们对True这里做无穷小  得到结果为0  

    dec_enc_attn_mask = get_attn_pad_mask(dec_inputs, enc_inputs) = 交互注意力机制PAD部分   其实这里应该是dec_inputs, enc_outputs  只不过不关注PAD用inputs足够了
            tensor([[[False, False, False, False,  True],                                 enc_outputs是生成阿巴阿巴 最后一个位置不是P 是P的新表现形式    中间这些都是未知的
                    [False, False, False, False,  True],
                    [False, False, False, False,  True],
                    [False, False, False, False,  True],
                    [False, False, False, False,  True]]])

    dec_self_attn_pad_mask = get_attn_pad_mask(dec_inputs, dec_inputs) =  dec_inputs自注意力层的时候PAD部分
            tensor([[[False, False, False, False, False],
                    [False, False, False, False, False],              横轴dec_inputs(S i want a beer) 和 竖轴 dec_inputs(S i want a beer)
                    [False, False, False, False, False],
                    [False, False, False, False, False],
                    [False, False, False, False, False]]])            
    '''

    '''Sequence Mask屏蔽子序列的mask '''

    @staticmethod
    def get_subsequent_mask(inputs):
        """
        屏蔽子序列的mask部分，这个函数就是用来表示Decoder的输入中哪些是未来词，使用一个上三角为1的矩阵遮蔽未来词，让当前词看不到未来词
        tensor([[[0, 1, 1, 1, 1],
                 [0, 0, 1, 1, 1],
                 [0, 0, 0, 1, 1],
                 [0, 0, 0, 0, 1],
                 [0, 0, 0, 0, 0]]], dtype=torch.uint8)
        :param inputs: shape(batch_size, max_len)
        """
        batch_size, max_len = inputs.size()
        attn_shape = [batch_size, max_len, max_len]
        subsequence_mask = np.triu(np.ones(attn_shape), k=1)
        # 生成一个上三角矩阵
        subsequence_mask = torch.from_numpy(subsequence_mask).byte()
        return subsequence_mask
